#!/usr/bin/env python

import numpy as N
import pyfits as pf
import sys

def nonlincor(file,*outfile):

  '''
  This function corrects NICMOS images for their count-rate dependent
  non-linearity. It used the header keywords CAMERA and FILTER to
  determine the non-linearity parameter. It corrects the first image,
  and in the case of a multi-extension image, the second image as well,
  with the appropriate power law.

  Checks still to be implemented:
   imagetyp
  '''

  # define the nonlinearity parameters
  wave1 = N.array([9034.6,11233.6,16036.9,21000,25000])
  nonl1 = N.array([0.101 ,0.074 ,0.031, 0.0, 0.0])
  
  wave2 = N.array([11234.7,16030.4,18705.6,21000,25000])
  nonl2 = N.array([0.063 ,0.029 ,0.013, 0.0, 0.0])

  wave3= N.array([.825,.875,.925,.975, 1.1, 1.2, 1.3, 1.4, 1.5, 1.6,  1.7,1.8,1.9,2.5])*10000
  nonl3 = N.array([.068,.056,.051,.048,.048,.048,.038,.026,.021,.016,.0105,.005,0,0])

  wave = (wave1,wave2,wave3)
  nonl = (nonl1,nonl2,nonl3)


  # open file and get camera and wavelength of filter
  name = file[0:min(file.index('.'),file.index('_'))]
  f = pf.open(file,)
  try:
    cam = int(f[0].header['CAMERA'])
  except KeyError:
    print("Required keyword CAMERA not found")
    raise KeyError
  try:
    filter = f[0].header['FILTER']
  except KeyError:
    print("Required keyword FILTER not found")
    raise KeyError


  # check whether the routine has already run on this data
  try:
    done =  f[0].header['NONLDONE']
    if (done == 'PERFORMED'):
      print("Non-linearity correction already been performed on %s"%(file))
      return()
  except KeyError:
    print("")

  # check whether we have Cycle 11+ data
  try:
    date =  f[0].header['DATE-OBS']
    if (int(date[0:4]) < 2000):
      # This will have to handle switch between Cycle 7 / 11+ eventually
      print("ERROR: This routine only works for Cycle 11+ data.")
      print("The calibration constants for Cycle 7 data have not yet been calculated.")
      return()
  except KeyError:
    print("WARNING: No DATE-OBS found. This routine only works for Cycle 11+ data.")

  # Warn about posibility of MultyDrizzle image being e-/s instead of ADU/s
  try:
    drz =  int(f[0].header['NDRIZIM'])
    print("WARNING: Detected multidrizzle image. Make sure image is in DN/s,")
    print("not electron/s. Otherwise, divide image by ADCGAIN value.")
  except KeyError:
    print("")

  # Warn about pedsky subtracting the sky
  try:
    skyval =  float(f[0].header['BKGND'])
    print("WARNING: Detected BKGND header keyword. This means that the sky may have been subtracted already. Add the BKGND into the image before using this routine.")
  except KeyError:
    print("")


  # get the photometric parameters for this camera/filter combination
  try:
    (photflam, plam, photfnu) = filterpar(cam,filter)
  except KeyError:
    print("Cannot handle this camera/filter combination")
    raise KeyError
  #print photflam, plam, photfnu

  # remember 0-indexed tulips
  icam = cam-1

  # calculate the non-linearty by linear interpolation
  ind = N.searchsorted(wave[icam],plam)-1
  nonlcor = nonl[icam][ind]+(plam-wave[icam][ind])*(nonl[icam][ind+1]-nonl[icam][ind])/(wave[icam][ind+1]-wave[icam][ind])

  print("Using non-linearity correction %6.4f mag/dex"%(nonlcor))
  
  # correct from mag/dex to alpha in power-law
  nonlcor /= 2.5
  nonlcor +=1
  #print nonlcor
  # now invert and subtract 1 to calculate the correction factor for each point
  inonlcor = (1./nonlcor)-1.
  #print nonlcor

  # calculate the new images
  if (len(f)==1):
    # simple fits file
    mul = N.where(N.not_equal(f[0].data,0),N.abs(f[0].data),1)**inonlcor
  if (len(f)>1):
    # multi extension fits file
    f[1].data += skyval
    mul = N.where(N.not_equal(f[1].data,0),N.abs(f[1].data),1)**inonlcor
    f[1].data *= mul
    f[2].data *= mul
    if abs(skyval)>0:
      nmul = N.abs(skyval)**inonlcor
      nskyval = skyval*nmul
      f[1].data -= nskyval

    # update header items in extensions
    f[1].header.update('NONLDONE','PERFORMED','corrected count-rate dependent non-linearity')
    f[1].header.update('NONLALPH',nonlcor,'power-law of non-linearity correction')
    f[1].header.update('PHOTFLAM',photflam,'inverse sensitivity (ergs/cm**2/Angstrom/DN)')
    f[1].header.update('PHOTFNU',photfnu,'inverse sensitivity (Jy*sec/DN)')
    f[2].header.update('NONLDONE','PERFORMED','corrected count-rate dependent non-linearity')
    f[2].header.update('NONLALPH',nonlcor,'power-law of non-linearity correction')
    f[2].header.update('PHOTFLAM',photflam,'inverse sensitivity (ergs/cm**2/Angstrom/DN)')
    f[2].header.update('PHOTFNU',photfnu,'inverse sensitivity (Jy*sec/DN)')

  # header items for top level
  f[0].header.update('NONLDONE','PERFORMED','corrected count-rate dependent non-linearity')
  f[0].header.update('NONLALPH',nonlcor,'power-law of non-linearity correction')
  f[0].header.update('PHOTFLAM',photflam,'inverse sensitivity (ergs/cm**2/Angstrom/DN)')
  f[0].header.update('PHOTFNU',photfnu,'inverse sensitivity (Jy*sec/DN)')

  # create output file
  if (outfile):
    f.writeto(outfile[0],clobber=True)
  else:
    f.writeto(name + '_nlc.fits',clobber=True)



def filterpar(camera,filter):

    filpar1 = {
    'F090M' :  (3.35572e-18,  9034.6, 9.13653e-06),
    'F095N' :  (4.58213e-17,  9536.8, 1.39012e-04),
    'F097N' :  (3.64484e-17,  9715.8, 1.14767e-04),
    'F108N' :  (2.48917e-17, 10816.3, 9.71381e-05),
    'F110M' :  (1.18916e-18, 11018.0, 4.81529e-06),
    'F110W' :  (4.05173e-19, 11233.6, 1.70552e-06),
    'F113N' :  (1.87273e-17, 11298.3, 7.97413e-05),
    'F140W' :  (1.45682e-19, 14268.5, 9.89330e-07),
    'F145M' :  (5.59605e-19, 14548.3, 3.95079e-06),
    'F160W' :  (2.08708e-19, 16036.9, 1.79044e-06),
    'F164N' :  (4.22990e-18, 16460.3, 3.82283e-05),
    'F165M' :  (3.82751e-19, 16480.3, 3.46757e-06),
    'F166N' :  (4.25553e-18, 16606.7, 3.91470e-05),
    'F170M' :  (3.37292e-19, 17057.9, 3.27367e-06),
    'F187N' :  (3.01786e-18, 18747.7, 3.53812e-05),
    'F190N' :  (2.96234e-18, 18985.9, 3.56186e-05)}

    filpar2 = {
    'F110W' :  (3.42604e-19, 11234.7, 1.44244e-06),
    'F160W' :  (1.87238e-19, 16030.4, 1.60496e-06),
    'F165M' :  (3.44504e-19, 16508.6, 3.13178e-06),
    'F171M' :  (8.07147e-19, 17211.0, 7.97524e-06),
    'F180M' :  (7.72258e-19, 17969.8, 8.31819e-06),
    'F187N' :  (2.72921e-18, 18739.8, 3.19703e-05),
    'F187W' :  (2.48423e-19, 18705.7, 2.89945e-06),
    'F190N' :  (2.71780e-18, 19003.4, 3.27385e-05),
    'F204M' :  (4.24793e-19, 20351.9, 5.86903e-06),
    'F205W' :  (6.85050e-20, 20636.1, 9.73097e-07),
    'F207M' :  (2.98464e-19, 20819.2, 4.31518e-06),
    'F212N' :  (1.91499e-18, 21212.8, 2.87436e-05),
    'F215N' :  (2.07570e-18, 21487.2, 3.19669e-05),
    'F216N' :  (1.86273e-18, 21640.9, 2.90991e-05),
    'F222M' :  (2.60915e-19, 22175.2, 4.27969e-06),
    'F237M' :  (1.98825e-19, 23690.5, 3.72217e-06)}

    filpar3 = {
    'F108N' :  (2.82793e-17, 10799.5, 1.10015e-04),
    'F110W' :  (4.38132e-19, 11200.5, 1.83340e-06),
    'F113N' :  (2.11233e-17, 11283.8, 8.97116e-05),
    'F150W' :  (1.32055e-19, 15349.8, 1.03787e-06),
    'F160W' :  (2.34763e-19, 16041.6, 2.01514e-06),
    'F164N' :  (4.74451e-18, 16460.3, 4.28791e-05),
    'F166N' :  (5.07451e-18, 16582.8, 4.65469e-05),
    'F175W' :  (6.54355e-20, 18090.6, 7.14327e-07),
    'F187N' :  (3.60365e-18, 18747.8, 4.22494e-05),
    'F190N' :  (3.33836e-18, 19003.4, 4.02138e-05),
    'F196N' :  (2.93916e-18, 19639.0, 3.78129e-05),
    'F200N' :  (2.75061e-18, 19974.8, 3.66075e-05),
    'F212N' :  (2.43465e-18, 21213.0, 3.65441e-05),
    'F215N' :  (2.60192e-18, 21487.2, 4.00714e-05),
    'F222M' :  (3.24648e-19, 22175.5, 5.32521e-06),
    'F240M' :  (1.96222e-19, 23957.4, 3.75672e-06)}

    filpar = {1:filpar1, 2:filpar2, 3:filpar3}    

    return filpar[camera][filter]


if __name__ == '__main__':
  len(sys.argv)
  if (len(sys.argv)==3):
    nonlincor(sys.argv[1],sys.argv[2])
  if (len(sys.argv)==2):
    nonlincor(sys.argv[1])
  if (len(sys.argv)==1):
    print("Usage:  nonlincor.py infile [outfile]")